#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
错误处理与重试机制模块
负责处理API调用过程中的各种异常情况，并提供智能重试策略
"""

import time
import random
import logging
from typing import Callable, Dict, Any, Optional, Union
import requests
from requests.exceptions import (
    RequestException,
    Timeout,
    ConnectionError,
    HTTPError,
    TooManyRedirects
)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('zhipu_api_error_handler')


class RetryStrategy:
    """
    重试策略类
    提供不同的重试策略实现
    """
    
    def __init__(self, max_retries: int = 3, base_delay: float = 1.0, max_delay: float = 60.0):
        """
        初始化重试策略
        
        Args:
            max_retries: 最大重试次数
            base_delay: 基础延迟时间（秒）
            max_delay: 最大延迟时间（秒）
        """
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
    
    def get_delay(self, attempt: int) -> float:
        """
        获取当前重试尝试的延迟时间
        
        Args:
            attempt: 当前重试次数（从1开始）
            
        Returns:
            float: 延迟时间（秒）
        """
        # 指数退避策略，加入随机抖动避免雪崩
        delay = self.base_delay * (2 ** (attempt - 1)) * (0.5 + random.random())
        return min(delay, self.max_delay)


class ApiErrorHandler:
    """
    API错误处理器
    负责处理API调用过程中的各种错误情况
    """
    
    # 需要重试的HTTP状态码
    RETRIABLE_STATUS_CODES = {
        429,  # Too Many Requests
        500,  # Internal Server Error
        502,  # Bad Gateway
        503,  # Service Unavailable
        504   # Gateway Timeout
    }
    
    # 错误码映射
    ERROR_CODE_MAP = {
        # 网络相关错误
        'network_error': '网络连接错误，请检查网络连接',
        'timeout_error': '请求超时，请稍后重试',
        'rate_limit_error': '请求频率超限，请降低请求频率',
        'server_error': '服务器错误，请稍后重试',
        
        # API相关错误
        'invalid_api_key': '无效的API密钥，请检查配置',
        'insufficient_quota': 'API配额不足，请充值或等待配额重置',
        'invalid_request': '无效的请求参数，请检查输入',
        'unauthorized': '未授权访问，请检查API密钥权限',
        
        # 其他错误
        'unknown_error': '未知错误，请稍后重试'
    }
    
    def __init__(self, retry_strategy: Optional[RetryStrategy] = None):
        """
        初始化错误处理器
        
        Args:
            retry_strategy: 重试策略实例，如果为None则使用默认策略
        """
        self.retry_strategy = retry_strategy or RetryStrategy()
        
    def should_retry(self, error: Exception) -> bool:
        """
        判断是否应该重试
        
        Args:
            error: 发生的异常
            
        Returns:
            bool: 是否应该重试
        """
        # 网络错误和超时错误应该重试
        if isinstance(error, (ConnectionError, Timeout)):
            return True
        
        # HTTP错误根据状态码判断
        if isinstance(error, HTTPError):
            return error.response.status_code in self.RETRIABLE_STATUS_CODES
        
        # 其他错误不重试
        return False
    
    def get_error_category(self, error: Exception) -> str:
        """
        获取错误类别
        
        Args:
            error: 发生的异常
            
        Returns:
            str: 错误类别
        """
        if isinstance(error, ConnectionError):
            return 'network_error'
        elif isinstance(error, Timeout):
            return 'timeout_error'
        elif isinstance(error, HTTPError):
            status_code = error.response.status_code
            if status_code == 429:
                return 'rate_limit_error'
            elif status_code >= 500:
                return 'server_error'
            elif status_code == 401:
                return 'unauthorized'
            elif status_code == 403:
                return 'invalid_api_key'
        
        return 'unknown_error'
    
    def get_error_message(self, error: Exception) -> str:
        """
        获取错误信息
        
        Args:
            error: 发生的异常
            
        Returns:
            str: 错误信息
        """
        category = self.get_error_category(error)
        base_message = self.ERROR_CODE_MAP.get(category, self.ERROR_CODE_MAP['unknown_error'])
        
        # 添加具体错误详情
        if hasattr(error, 'response') and error.response is not None:
            try:
                error_data = error.response.json()
                if 'error' in error_data:
                    if isinstance(error_data['error'], dict):
                        return f"{base_message}: {error_data['error'].get('message', str(error_data['error']))}"
                    return f"{base_message}: {error_data['error']}"
            except (ValueError, KeyError):
                pass
        
        # 添加异常的原始消息
        if str(error):
            return f"{base_message}: {str(error)}"
        
        return base_message
    
    def handle_api_call(self, func: Callable, *args, **kwargs) -> Dict[str, Any]:
        """
        处理API调用，包含重试逻辑
        
        Args:
            func: 要调用的函数
            *args: 函数参数
            **kwargs: 关键字参数
            
        Returns:
            Dict[str, Any]: 包含调用结果的字典
        """
        last_error = None
        
        for attempt in range(self.retry_strategy.max_retries + 1):
            try:
                # 如果不是第一次尝试，记录重试信息
                if attempt > 0:
                    delay = self.retry_strategy.get_delay(attempt)
                    logger.warning(
                        f"API调用失败，正在进行第{attempt}/{self.retry_strategy.max_retries}次重试，" 
                        f"等待{delay:.2f}秒..."
                    )
                    time.sleep(delay)
                
                # 调用函数
                logger.info(f"执行API调用 (尝试 {attempt+1}/{self.retry_strategy.max_retries+1})")
                result = func(*args, **kwargs)
                
                # 检查结果是否成功
                if isinstance(result, dict) and result.get('success', True):
                    logger.info(f"API调用成功 (尝试 {attempt+1}/{self.retry_strategy.max_retries+1})")
                    return result
                elif isinstance(result, dict) and 'error' in result:
                    last_error = Exception(result['error'])
                    logger.error(f"API返回错误: {result['error']}")
                else:
                    last_error = Exception(f"API返回意外结果: {str(result)}")
                    logger.error(f"API返回意外结果: {str(result)}")
                    
            except Exception as e:
                last_error = e
                logger.error(f"API调用异常 (尝试 {attempt+1}/{self.retry_strategy.max_retries+1}): {str(e)}")
                
                # 判断是否应该重试
                if not self.should_retry(e):
                    logger.warning(f"遇到不可重试的错误，停止重试")
                    break
            
            # 如果是最后一次尝试，不再重试
            if attempt == self.retry_strategy.max_retries:
                break
        
        # 所有尝试都失败了
        error_message = self.get_error_message(last_error)
        logger.error(f"所有重试都失败了: {error_message}")
        
        return {
            'success': False,
            'error': error_message,
            'attempts': self.retry_strategy.max_retries + 1,
            'last_error': str(last_error)
        }


class ZhipuApiErrorHandler(ApiErrorHandler):
    """
    智谱API专用错误处理器
    针对智谱API的特定错误类型进行处理
    """
    
    # 智谱API特定错误码
    ZHIPU_ERROR_CODES = {
        1001: '请求参数错误',
        1002: '未找到指定模型',
        1003: 'API密钥错误',
        1004: 'API密钥过期',
        1005: '账户余额不足',
        1006: '并发请求超限',
        1007: '请求频率超限',
        1008: '请求超时',
        1009: '服务器内部错误',
        1010: '模型生成内容被拦截',
        1011: '模型处理失败',
        1012: '文件大小超限',
        1013: '文件格式不支持',
        1014: '文件下载失败'
    }
    
    def should_retry(self, error: Exception) -> bool:
        """
        判断是否应该重试（针对智谱API的特定逻辑）
        
        Args:
            error: 发生的异常
            
        Returns:
            bool: 是否应该重试
        """
        # 先使用父类的判断逻辑
        if super().should_retry(error):
            return True
        
        # 检查智谱特定错误
        if isinstance(error, HTTPError):
            try:
                error_data = error.response.json()
                error_code = error_data.get('error', {}).get('code')
                
                # 这些错误码可以重试
                retryable_codes = {1008, 1009, 1011, 1014}
                if error_code in retryable_codes:
                    return True
                    
                # 频率限制相关的错误也可以重试
                if error_code in {1006, 1007}:
                    return True
                    
            except (ValueError, KeyError):
                pass
        
        return False
    
    def get_error_message(self, error: Exception) -> str:
        """
        获取智谱API特定的错误信息
        
        Args:
            error: 发生的异常
            
        Returns:
            str: 错误信息
        """
        # 尝试从响应中提取智谱API特定的错误信息
        if isinstance(error, HTTPError):
            try:
                error_data = error.response.json()
                if 'error' in error_data:
                    error_info = error_data['error']
                    if isinstance(error_info, dict):
                        error_code = error_info.get('code')
                        if error_code in self.ZHIPU_ERROR_CODES:
                            return f"{self.ZHIPU_ERROR_CODES[error_code]} (错误码: {error_code})"
                        return f"智谱API错误: {error_info.get('message', str(error_info))}"
                    return f"智谱API错误: {error_info}"
            except (ValueError, KeyError):
                pass
        
        # 如果无法提取智谱特定错误，使用通用错误信息
        return super().get_error_message(error)